from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

setup(
    name='step2',
    ext_modules=[
        CUDAExtension(
            name='ops_test_op',
            sources=[
                'ops_test.cpp',
                'ops_test_cuda.cu',#不能重名 否则会报多重定义的问题
                     ],
            define_macros=[('WITH_CUDA', None)]
        ),
    ],
    cmdclass={
        'build_ext': BuildExtension
    })
